import numpy as np
import matplotlib.pyplot as plt

# 物理常量和初始条件定义
g = 9.8  # 重力加速度 (m/s²)
Ct = np.array([0, 200, 0])  # 目标位置坐标 (x, y, z)
R_cloud = 10  # 遮蔽云团的半径 (m)
effective_duration = 20  # 云团有效作用时间 (s)

# 导弹初始状态
M1_start = np.array([20000, 0, 2000])  # 导弹初始位置 (x, y, z)
D_m = Ct - M1_start  # 导弹到目标的向量
dist_m = np.linalg.norm(D_m)  # 导弹到目标的直线距离
U_m = D_m / dist_m  # 导弹飞行方向单位向量
V_m = 300 * U_m  # 导弹速度向量 (m/s)，速率300m/s

# 无人机初始状态
FY1_start = np.array([17800, 0, 1800])  # 无人机初始位置 (x, y, z)


class LyapunovGuidedPSO:
    """结合李雅普诺夫(Lyapunov)引导策略的粒子群优化算法类"""
    
    def __init__(self, n_particles, dimensions, bounds, w=0.5, c1=1.5, c2=1.5):
        """
        初始化PSO参数
        
        参数:
            n_particles: 粒子数量
            dimensions: 优化问题的维度
            bounds: 每个维度的参数边界，格式为[(min1, max1), (min2, max2), ...]
            w: 惯性权重，控制粒子速度的继承
            c1: 认知学习因子，控制粒子向个体最优学习的权重
            c2: 社会学习因子，控制粒子向全局最优学习的权重
        """
        self.n_particles = n_particles  # 粒子数量
        self.dimensions = dimensions  # 优化维度
        self.bounds = bounds  # 参数边界
        self.w = w  # 惯性权重
        self.c1 = c1  # 认知因子
        self.c2 = c2  # 社会因子
        
        # 初始化粒子位置（在边界范围内随机生成）
        self.positions = np.random.uniform(
            low=[b[0] for b in bounds],  # 各维度下界
            high=[b[1] for b in bounds],  # 各维度上界
            size=(n_particles, dimensions)  # 生成(n_particles×dimensions)的位置矩阵
        )
        
        self.velocities = np.zeros((n_particles, dimensions))  # 初始化粒子速度为0
        
        # 初始化个体最优（初始为粒子自身位置）
        self.pbest_positions = self.positions.copy()  # 个体最优位置
        self.pbest_fitness = np.full(n_particles, np.inf)  # 个体最优适应度（初始为无穷大）
        
        # 初始化全局最优
        self.gbest_position = None  # 全局最优位置
        self.gbest_fitness = np.inf  # 全局最优适应度（初始为无穷大）
        
        # 李雅普诺夫函数相关变量
        self.V = np.zeros(n_particles)  # 李雅普诺夫函数值
        self.dV = np.zeros(n_particles)  # 李雅普诺夫函数变化率
        self.V_prev = np.zeros(n_particles)  # 上一时刻李雅普诺夫函数值
    
    def evaluate_fitness(self, position):
        """
        评估粒子位置的适应度（目标：最大化云团对导弹的遮蔽时间）
        
        参数:
            position: 粒子位置，包含4个参数[v, phi, t_drop, delta_t]
            
        返回:
            负的最大遮蔽时间（因PSO默认最小化适应度，故取负值）
        """
        # 解析粒子位置参数
        v, phi, t_drop, delta_t = position  # v:无人机速度；phi:方位角；t_drop:投放时间；delta_t:起爆延迟
        
        # 参数范围检查（超出范围则返回惩罚值）
        if (v < 70 or v > 140 or  # 速度范围限制：70-140m/s
            t_drop < 0 or delta_t < 0 or  # 时间不能为负
            phi < 0 or phi > 2*np.pi):  # 方位角范围：0到2π
            return 1e10  # 惩罚值（很大的正数，使该粒子不被选中）
        
        # 计算无人机飞行方向单位向量（水平方向，无仰角）
        ux = np.cos(phi)  # x方向分量
        uy = np.sin(phi)  # y方向分量
        uz = 0  # z方向分量为0（水平飞行）
        U_fy1 = np.array([ux, uy, uz])  # 无人机飞行方向单位向量
        V_fy1 = v * U_fy1  # 无人机速度向量（速率×方向）

        # 计算投放点和起爆点位置
        P_drop = FY1_start + V_fy1 * t_drop  # 投放点位置 = 初始位置 + 速度×投放时间
        t_burst = t_drop + delta_t  # 起爆时间 = 投放时间 + 延迟时间
        V_drop = V_fy1  # 投放物初速度 = 无人机速度（假设投放时速度相同）
        # 起爆点位置：投放点 + 投放物运动（含重力影响）
        P_burst = P_drop + V_drop * delta_t + np.array([0, 0, -0.5 * g * delta_t**2])

        # 模拟云团有效时间内的遮蔽效果
        t_min = t_burst  # 云团生效起始时间（起爆时间）
        t_max = t_burst + effective_duration  # 云团生效结束时间
        t_values = np.arange(t_min, t_max, 0.1)  # 时间采样点（步长0.1s）
        obscured = []  # 记录每个时刻是否遮蔽
        
        for t in t_values:
            C_c = self.cloud_center(t_burst, P_burst, t)  # 计算t时刻云团中心位置
            is_obscured = self.calculate_obscuration(t, C_c)  # 判断该时刻是否遮蔽
            obscured.append(is_obscured)
        
        # 计算最长连续遮蔽时间
        obscured = np.array(obscured)  # 转换为数组便于处理
        changes = np.diff(obscured.astype(int))  # 计算状态变化（0→1或1→0）
        start_indices = np.where(changes == 1)[0] + 1  # 遮蔽开始的索引
        end_indices = np.where(changes == -1)[0] + 1  # 遮蔽结束的索引
        
        # 处理边界情况（若起始时刻已遮蔽）
        if obscured[0]:
            start_indices = np.insert(start_indices, 0, 0)
        # 处理边界情况（若结束时刻仍遮蔽）
        if obscured[-1]:
            end_indices = np.append(end_indices, len(obscured))
        
        max_duration = 0  # 最长遮蔽时间
        for start, end in zip(start_indices, end_indices):
            duration = t_values[end] - t_values[start]  # 计算当前遮蔽段的时长
            if duration > max_duration:
                max_duration = duration  # 更新最长时长
        
        return -max_duration  # 返回负的最长时长（用于最小化优化）
    
    def cloud_center(self, t_burst, P_burst, t):
        """
        计算云团中心在t时刻的位置（简化模型）
        
        参数:
            t_burst: 起爆时间
            P_burst: 起爆点位置
            t: 当前时刻
            
        返回:
            云团中心坐标 (x, y, z)
        """
        # 假设云团在x、y方向不动，z方向（高度）随时间下降（3m/s）
        z = P_burst[2] - 3 * (t - t_burst)
        return np.array([P_burst[0], P_burst[1], z])
    
    def calculate_obscuration(self, t, C_cloud):
        """
        判断t时刻云团是否遮蔽导弹到目标的视线
        
        参数:
            t: 当前时刻
            C_cloud: 云团中心位置
            
        返回:
            bool: 若遮蔽则为True，否则为False
        """
        P_m = M1_start + V_m * t  # 计算t时刻导弹位置
        
        # 计算导弹到目标的视线向量
        target_vector = Ct - P_m  # 目标 - 导弹位置向量
        target_dist = np.linalg.norm(target_vector)  # 导弹到目标的距离
        target_dir = target_vector / target_dist  # 视线方向单位向量
        
        # 计算云团中心到视线的最短距离（几何判断）
        cloud_to_line = C_cloud - P_m  # 云团中心到导弹位置的向量
        projection = np.dot(cloud_to_line, target_dir)  # 向量在视线上的投影长度
        
        # 根据投影位置计算最短距离
        if projection < 0:
            # 投影在导弹后方→最短距离为云团到导弹的距离
            dist = np.linalg.norm(cloud_to_line)
        elif projection > target_dist:
            # 投影在目标前方→最短距离为云团到目标的距离
            dist = np.linalg.norm(C_cloud - Ct)
        else:
            # 投影在视线中间→计算垂线距离
            closest_point = P_m + projection * target_dir  # 垂点位置
            dist = np.linalg.norm(C_cloud - closest_point)  # 云团到垂点的距离
        
        # 若最短距离≤云团半径，则判定为遮蔽
        return dist <= R_cloud
    
    def update_lyapunov(self, iteration):
        """
        更新李雅普诺夫函数值及变化率（用于引导粒子探索）
        
        参数:
            iteration: 当前迭代次数
        """
        for i in range(self.n_particles):
            # 李雅普诺夫函数定义为个体最优与全局最优的适应度差的平方
            fitness_diff = self.pbest_fitness[i] - self.gbest_fitness
            self.V[i] = fitness_diff**2
            
            # 计算李雅普诺夫函数变化率（当前值 - 上一时刻值）
            if iteration > 0:
                self.dV[i] = self.V[i] - self.V_prev[i]
            else:
                self.dV[i] = 0  # 初始时刻变化率为0
        
        self.V_prev = self.V.copy()  # 保存当前值用于下一时刻计算
    
    def update_particles(self, iteration):
        """
        更新粒子的速度和位置，并结合李雅普诺夫引导策略增强探索
        
        参数:
            iteration: 当前迭代次数
        """
        self.update_lyapunov(iteration)  # 先更新李雅普诺夫函数
        
        for i in range(self.n_particles):
            # 生成随机因子（0-1之间）
            r1 = np.random.random(self.dimensions)  # 认知项随机因子
            r2 = np.random.random(self.dimensions)  # 社会项随机因子
            
            # 计算速度更新量
            cognitive = self.c1 * r1 * (self.pbest_positions[i] - self.positions[i])  # 认知项（向个体最优学习）
            social = self.c2 * r2 * (self.gbest_position - self.positions[i])  # 社会项（向全局最优学习）
            # 速度公式：惯性×当前速度 + 认知项 + 社会项
            self.velocities[i] = self.w * self.velocities[i] + cognitive + social
            
            # 李雅普诺夫引导的探索增强：当dV<0时（函数值下降），增加随机扰动
            if self.dV[i] < 0:
                exploration_factor = 0.1 * (1 - iteration / 100)  # 探索因子随迭代衰减
                random_exploration = exploration_factor * np.random.randn(self.dimensions)  # 高斯分布随机扰动
                self.velocities[i] += random_exploration  # 增强探索
            
            # 更新粒子位置
            self.positions[i] += self.velocities[i]
            
            # 边界处理：若超出边界，将位置拉回边界并反向速度（部分反弹）
            for d in range(self.dimensions):
                if self.positions[i, d] < self.bounds[d][0]:
                    self.positions[i, d] = self.bounds[d][0]  # 限制在左边界
                    self.velocities[i, d] *= -0.5  # 速度反向并衰减
                elif self.positions[i, d] > self.bounds[d][1]:
                    self.positions[i, d] = self.bounds[d][1]  # 限制在右边界
                    self.velocities[i, d] *= -0.5  # 速度反向并衰减
    
    def optimize(self, max_iterations):
        """
        执行PSO优化过程
        
        参数:
            max_iterations: 最大迭代次数
            
        返回:
            gbest_position: 全局最优位置（最优参数）
            -gbest_fitness: 最大遮蔽时间（转换为正值）
            convergence_curve: 收敛曲线（记录每代最优适应度）
        """
        convergence_curve = []  # 记录收敛过程
        
        # 初始评估（第一代粒子）
        for i in range(self.n_particles):
            fitness = self.evaluate_fitness(self.positions[i])  # 计算适应度
            
            # 更新个体最优
            if fitness < self.pbest_fitness[i]:
                self.pbest_fitness[i] = fitness
                self.pbest_positions[i] = self.positions[i].copy()
            
            # 更新全局最优
            if fitness < self.gbest_fitness:
                self.gbest_fitness = fitness
                self.gbest_position = self.positions[i].copy()
        
        self.update_lyapunov(0)  # 初始化李雅普诺夫函数
        convergence_curve.append(self.gbest_fitness)  # 记录初始最优
        
        # 迭代优化
        for iteration in range(1, max_iterations):
            self.update_particles(iteration)  # 更新粒子位置和速度
            
            # 评估并更新最优解
            for i in range(self.n_particles):
                fitness = self.evaluate_fitness(self.positions[i])  # 计算当前适应度
                
                # 更新个体最优
                if fitness < self.pbest_fitness[i]:
                    self.pbest_fitness[i] = fitness
                    self.pbest_positions[i] = self.positions[i].copy()
                
                # 更新全局最优
                if fitness < self.gbest_fitness:
                    self.gbest_fitness = fitness
                    self.gbest_position = self.positions[i].copy()
            
            convergence_curve.append(self.gbest_fitness)  # 记录当前代最优
            
            # 每10代打印一次进度
            if iteration % 10 == 0:
                print(f"Iteration {iteration}, Best Fitness: {-self.gbest_fitness:.4f}")
        
        return self.gbest_position, -self.gbest_fitness, convergence_curve  # 转换适应度为正值（遮蔽时间）


def problem2_lyapunov_pso():
    """使用李雅普诺夫引导的PSO算法求解遮蔽优化问题"""
    # 定义各参数的边界范围
    bounds = [
        (70, 140),          # 无人机速度 v (m/s)
        (0, 2*np.pi),       # 方位角 phi (rad) - 0到2π，X轴正方向为0，逆时针为正
        (0, 10),            # 投放时间 t_drop (s)
        (0, 10)             # 起爆延迟 delta_t (s)
    ]
    
    # 创建PSO优化器实例
    pso = LyapunovGuidedPSO(
        n_particles=30,    # 30个粒子
        dimensions=4,      # 4个优化参数（速度、方位角、投放时间、起爆延迟）
        bounds=bounds,     # 参数边界
        w=0.7,             # 惯性权重
        c1=1.5,            # 认知因子
        c2=1.5             # 社会因子
    )
    
    # 执行优化（最大迭代100次）
    best_position, best_fitness, convergence_curve = pso.optimize(max_iterations=100)
    
    # 提取最优参数
    v_opt, phi_opt, t_drop_opt, delta_t_opt = best_position
    t_burst_opt = t_drop_opt + delta_t_opt  # 最优起爆时间
    
    # 将方位角从弧度转换为角度（0-360度）
    phi_deg = np.degrees(phi_opt) % 360
    
    # 输出优化结果
    print("\n=== 优化结果 ===")
    print(f"无人机速度: {v_opt:.2f} m/s")
    print(f"方位角: {phi_deg:.2f}° (从X轴正方向逆时针测量)")
    print(f"投放时间: {t_drop_opt:.2f} s")
    print(f"起爆延迟: {delta_t_opt:.2f} s")
    print(f"起爆时间: {t_burst_opt:.2f} s")
    print(f"最大遮蔽时长: {best_fitness:.2f}秒")
    
    # 绘制收敛曲线
    plt.figure(figsize=(10, 6))
    plt.plot(convergence_curve)
    plt.title('PSO收敛曲线')
    plt.xlabel('迭代次数')
    plt.ylabel('适应度值 (负的遮蔽时间)')
    plt.grid(True)
    plt.savefig('convergence_curve_p2.png')  # 保存图像
    plt.show()
    
    return best_position, best_fitness


# 主程序入口
if __name__ == "__main__":
    problem2_lyapunov_pso()  # 运行优化函数